#!/bin/bash

# Flink HBase ETL 作业提交脚本
# 作者: Flink HBase ETL 团队
# 版本: 1.0.0

set -e

# 颜色定义
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
BLUE='\033[0;34m'
NC='\033[0m' # No Color

# 日志函数
log_info() {
    echo -e "${BLUE}[INFO]${NC} $1"
}

log_warn() {
    echo -e "${YELLOW}[WARN]${NC} $1"
}

log_error() {
    echo -e "${RED}[ERROR]${NC} $1"
}

log_success() {
    echo -e "${GREEN}[SUCCESS]${NC} $1"
}

# 打印使用说明
print_usage() {
    cat << EOF
Flink HBase ETL 作业提交脚本

使用方法:
    $0 [选项] <table_name> <output_path>

位置参数:
    table_name      HBase 表名
    output_path     HDFS 输出路径

选项:
    -h, --help              显示此帮助信息
    -e, --env ENV           环境 (dev|test|prod) [默认: dev]
    -p, --parallelism N     并行度 [默认: 4]
    -j, --jar JAR_PATH      JAR 文件路径 [默认: target/flink-hbase-etl-1.0-SNAPSHOT.jar]
    -m, --mode MODE         运行模式 (local|yarn-session|yarn-per-job) [默认: yarn-per-job]
    -s, --source TYPE       数据源类型 (mock|file|kafka) [默认: mock]
    -f, --file FILE_PATH    输入文件路径 (当 source=file 时)
    -k, --kerberos          启用 Kerberos 认证
    -c, --checkpoint MS     检查点间隔毫秒数 [默认: 60000]
    -d, --dry-run           只显示命令，不实际执行

示例:
    # 本地开发环境测试
    $0 -e dev -p 2 -s mock user /tmp/output

    # 生产环境运行
    $0 -e prod -p 8 -s file -f /data/input/user_ids.txt -k user /data/warehouse/user/

    # YARN Session 模式
    $0 -m yarn-session -p 4 user /data/warehouse/user/

EOF
}

# 默认参数
ENV="dev"
PARALLELISM="4"
JAR_PATH="target/flink-hbase-etl-1.0-SNAPSHOT.jar"
MODE="yarn-per-job"
SOURCE_TYPE="mock"
INPUT_FILE=""
ENABLE_KERBEROS=false
CHECKPOINT_INTERVAL="60000"
DRY_RUN=false

# 解析命令行参数
while [[ $# -gt 0 ]]; do
    case $1 in
        -h|--help)
            print_usage
            exit 0
            ;;
        -e|--env)
            ENV="$2"
            shift 2
            ;;
        -p|--parallelism)
            PARALLELISM="$2"
            shift 2
            ;;
        -j|--jar)
            JAR_PATH="$2"
            shift 2
            ;;
        -m|--mode)
            MODE="$2"
            shift 2
            ;;
        -s|--source)
            SOURCE_TYPE="$2"
            shift 2
            ;;
        -f|--file)
            INPUT_FILE="$2"
            shift 2
            ;;
        -k|--kerberos)
            ENABLE_KERBEROS=true
            shift
            ;;
        -c|--checkpoint)
            CHECKPOINT_INTERVAL="$2"
            shift 2
            ;;
        -d|--dry-run)
            DRY_RUN=true
            shift
            ;;
        -*)
            log_error "未知选项: $1"
            print_usage
            exit 1
            ;;
        *)
            break
            ;;
    esac
done

# 检查位置参数
if [ $# -lt 2 ]; then
    log_error "缺少必需参数: table_name 和 output_path"
    print_usage
    exit 1
fi

TABLE_NAME="$1"
OUTPUT_PATH="$2"

# 验证环境
case $ENV in
    dev|test|prod)
        ;;
    *)
        log_error "无效的环境: $ENV。支持的环境: dev, test, prod"
        exit 1
        ;;
esac

# 验证运行模式
case $MODE in
    local|yarn-session|yarn-per-job)
        ;;
    *)
        log_error "无效的运行模式: $MODE。支持的模式: local, yarn-session, yarn-per-job"
        exit 1
        ;;
esac

# 验证数据源类型
case $SOURCE_TYPE in
    mock|file|kafka)
        ;;
    *)
        log_error "无效的数据源类型: $SOURCE_TYPE。支持的类型: mock, file, kafka"
        exit 1
        ;;
esac

# 如果是文件数据源，检查文件路径
if [ "$SOURCE_TYPE" = "file" ] && [ -z "$INPUT_FILE" ]; then
    log_error "使用文件数据源时必须指定 --file 参数"
    exit 1
fi

# 检查 JAR 文件存在
if [ ! -f "$JAR_PATH" ]; then
    log_error "JAR 文件不存在: $JAR_PATH"
    log_info "请先运行: mvn clean package"
    exit 1
fi

# 环境变量配置
export HADOOP_CONF_DIR=${HADOOP_CONF_DIR:-/etc/hadoop/conf}
export HBASE_CONF_DIR=${HBASE_CONF_DIR:-/etc/hbase/conf}
export FLINK_HOME=${FLINK_HOME:-/opt/flink}

# 根据环境设置配置
case $ENV in
    dev)
        JM_MEMORY="1024m"
        TM_MEMORY="2048m"
        ASYNC_POOL_SIZE="5"
        FILE_SIZE_MB="64"
        ;;
    test)
        JM_MEMORY="2048m"
        TM_MEMORY="4096m"
        ASYNC_POOL_SIZE="10"
        FILE_SIZE_MB="128"
        ;;
    prod)
        JM_MEMORY="4096m"
        TM_MEMORY="8192m"
        ASYNC_POOL_SIZE="20"
        FILE_SIZE_MB="256"
        ;;
esac

# 构建 Flink 命令
if [ "$MODE" = "local" ]; then
    FLINK_CMD="java -cp $JAR_PATH com.flink.hbase.JobLauncher"
else
    FLINK_CMD="$FLINK_HOME/bin/flink run"
    
    # 添加 YARN 参数
    if [[ "$MODE" == yarn-* ]]; then
        FLINK_CMD="$FLINK_CMD --target $MODE"
        FLINK_CMD="$FLINK_CMD --jobmanager-memory $JM_MEMORY"
        FLINK_CMD="$FLINK_CMD --taskmanager-memory $TM_MEMORY"
        
        # Kerberos 文件
        if [ "$ENABLE_KERBEROS" = true ]; then
            KERBEROS_FILES="hdfs:///user/flink/conf/krb5.conf,hdfs:///user/flink/conf/user.keytab"
            HADOOP_CONFIG_FILES="hdfs:///user/flink/conf/core-site.xml,hdfs:///user/flink/conf/hdfs-site.xml,hdfs:///user/flink/conf/hbase-site.xml"
            FLINK_CMD="$FLINK_CMD --files \"$KERBEROS_FILES,$HADOOP_CONFIG_FILES\""
        fi
    fi
    
    FLINK_CMD="$FLINK_CMD --parallelism $PARALLELISM"
    FLINK_CMD="$FLINK_CMD $JAR_PATH"
fi

# 构建应用参数
APP_ARGS=""
APP_ARGS="$APP_ARGS --job.name \"HBase-ETL-$ENV-$(date +%Y%m%d-%H%M%S)\""
APP_ARGS="$APP_ARGS --hbase.table.name $TABLE_NAME"
APP_ARGS="$APP_ARGS --output.path $OUTPUT_PATH"
APP_ARGS="$APP_ARGS --source.type $SOURCE_TYPE"
APP_ARGS="$APP_ARGS --checkpoint.interval $CHECKPOINT_INTERVAL"
APP_ARGS="$APP_ARGS --async.pool.size $ASYNC_POOL_SIZE"
APP_ARGS="$APP_ARGS --sink.file.size.mb $FILE_SIZE_MB"

# 数据源特定参数
case $SOURCE_TYPE in
    file)
        APP_ARGS="$APP_ARGS --source.file.path $INPUT_FILE"
        ;;
    mock)
        MOCK_DATA_SIZE="1000"
        if [ "$ENV" = "prod" ]; then
            MOCK_DATA_SIZE="10000"
        fi
        APP_ARGS="$APP_ARGS --mock.data.size $MOCK_DATA_SIZE"
        ;;
esac

# Kerberos 参数
if [ "$ENABLE_KERBEROS" = true ]; then
    APP_ARGS="$APP_ARGS --security.kerberos.enabled true"
    APP_ARGS="$APP_ARGS --security.kerberos.login.principal user@REALM.COM"
    APP_ARGS="$APP_ARGS --security.kerberos.login.keytab ./user.keytab"
fi

# 环境特定的 HBase 配置
case $ENV in
    dev)
        APP_ARGS="$APP_ARGS --hbase.zookeeper.quorum localhost"
        ;;
    test)
        APP_ARGS="$APP_ARGS --hbase.zookeeper.quorum test-zk1,test-zk2,test-zk3"
        ;;
    prod)
        APP_ARGS="$APP_ARGS --hbase.zookeeper.quorum prod-zk1,prod-zk2,prod-zk3"
        ;;
esac

# 完整命令
FULL_CMD="$FLINK_CMD $APP_ARGS"

# 显示配置信息
log_info "========== 作业配置 =========="
log_info "环境: $ENV"
log_info "运行模式: $MODE"
log_info "HBase 表名: $TABLE_NAME"
log_info "输出路径: $OUTPUT_PATH"
log_info "数据源: $SOURCE_TYPE"
log_info "并行度: $PARALLELISM"
log_info "Kerberos: $([ "$ENABLE_KERBEROS" = true ] && echo "启用" || echo "禁用")"
log_info "检查点间隔: ${CHECKPOINT_INTERVAL}ms"
log_info "================================"

# 显示完整命令
log_info "执行命令:"
echo "$FULL_CMD"
echo

# 确认执行
if [ "$DRY_RUN" = true ]; then
    log_warn "这是一次试运行，不会实际执行作业"
    exit 0
fi

read -p "是否确认提交作业? (y/N): " confirm
if [[ ! $confirm =~ ^[Yy]$ ]]; then
    log_warn "作业提交已取消"
    exit 0
fi

# 执行命令
log_info "正在提交作业..."
eval $FULL_CMD

if [ $? -eq 0 ]; then
    log_success "作业提交成功!"
    log_info "作业名称: HBase-ETL-$ENV-$(date +%Y%m%d-%H%M%S)"
    log_info "可以通过 Flink Web UI 查看作业状态"
else
    log_error "作业提交失败"
    exit 1
fi 